package com.mm.cloth.rmbg;

import ai.onnxruntime.*;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.Collections;

/**
 * 简化版的LaMa推理助手
 * 专门针对200MB big-lama_fp16.onnx模型优化
 * 使用最简单的配置避免卡住
 */
public class SimpleLamaInferenceHelper {
    private static final String TAG = "SimpleLamaInference";
    
    private OrtSession session;
    private OrtEnvironment env;
    
    // 模型参数
    private static final String INPUT_NAME = "input";
    private static final String OUTPUT_NAME = "output";
    
    /**
     * 初始化模型（简化版本）
     */
    public void initModel(AssetManager assets, String modelPath) throws Exception {
        Log.d(TAG, "=== 简化版Big-LaMa标准模型初始化开始 ===");
        Log.d(TAG, "模型路径: " + modelPath);
        
        try {
            // 1. 检查ONNX Runtime
            try {
                Class.forName("ai.onnxruntime.OrtEnvironment");
                Log.d(TAG, "✅ ONNX Runtime可用");
            } catch (ClassNotFoundException e) {
                throw new Exception("ONNX Runtime不可用", e);
            }
            
            // 2. 创建环境
            Log.d(TAG, "创建ONNX Runtime环境...");
            env = OrtEnvironment.getEnvironment();
            Log.d(TAG, "✅ ONNX Runtime环境创建成功");
            
            // 3. 检查内存
            checkMemory();
            
            // 4. 加载模型文件
            Log.d(TAG, "开始加载模型文件...");
            byte[] modelBytes = loadModelFile(assets, modelPath);
            Log.d(TAG, "✅ 模型文件加载完成: " + (modelBytes.length / (1024*1024)) + "MB");
            
            // 5. 创建最简单的会话选项
            Log.d(TAG, "配置会话选项（最简化）...");
            OrtSession.SessionOptions options = new OrtSession.SessionOptions();
            
            // 使用最保守的设置
            options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT);
            options.setIntraOpNumThreads(1);
            options.setInterOpNumThreads(1);
            
            Log.d(TAG, "✅ 会话选项配置完成（保守模式）");
            
            // 6. 强制垃圾回收
            Log.d(TAG, "执行垃圾回收...");
            System.gc();
            Thread.sleep(500); // 给GC更多时间
            
            // 7. 创建会话
            Log.d(TAG, "开始创建ONNX会话（可能需要10-30秒）...");
            long startTime = System.currentTimeMillis();
            
            session = env.createSession(modelBytes, options);
            
            long endTime = System.currentTimeMillis();
            Log.d(TAG, "✅ ONNX会话创建成功！耗时: " + (endTime - startTime) + "ms");
            
            // 8. 验证会话
            if (session == null) {
                throw new Exception("会话创建返回null");
            }
            
            // 9. 打印模型信息
            printModelInfo();
            
            Log.d(TAG, "=== 简化版Big-LaMa模型初始化完成 ===");
            
        } catch (Exception e) {
            Log.e(TAG, "模型初始化失败", e);
            cleanup();
            throw e;
        }
    }
    
    /**
     * 检查内存状态
     */
    private void checkMemory() {
        Runtime runtime = Runtime.getRuntime();
        long maxMemory = runtime.maxMemory();
        long totalMemory = runtime.totalMemory();
        long freeMemory = runtime.freeMemory();
        long usedMemory = totalMemory - freeMemory;
        long availableMemory = maxMemory - usedMemory;
        
        long availableMB = availableMemory / (1024 * 1024);
        long maxMB = maxMemory / (1024 * 1024);
        
        Log.d(TAG, "内存状态 - 最大: " + maxMB + "MB, 可用: " + availableMB + "MB");
        
        if (availableMB < 120) {
            Log.w(TAG, "⚠️ 可用内存不足120MB，可能无法加载Big-LaMa模型");
            throw new RuntimeException("内存不足，可用: " + availableMB + "MB，需要至少120MB");
        } else {
            Log.d(TAG, "✅ 内存充足，可以加载大模型");
        }
    }
    
    /**
     * 加载模型文件
     */
    private byte[] loadModelFile(AssetManager assets, String modelPath) throws Exception {
        try (InputStream modelStream = assets.open(modelPath)) {
            int available = modelStream.available();
            float sizeMB = available / (1024.0f * 1024.0f);
            
            Log.d(TAG, "模型文件大小: " + String.format("%.1f", sizeMB) + "MB");
            
            if (available == 0) {
                throw new Exception("模型文件为空");
            }
            
            if (available < 50 * 1024 * 1024) { // 小于50MB
                throw new Exception("模型文件太小，可能不是big-lama.onnx");
            }
            
            // 简单一次性读取（对于200MB文件）
            byte[] modelBytes = new byte[available];
            int totalRead = 0;
            
            while (totalRead < available) {
                int read = modelStream.read(modelBytes, totalRead, available - totalRead);
                if (read == -1) break;
                totalRead += read;
                
                // 每50MB显示一次进度
                if (totalRead % (50 * 1024 * 1024) == 0 || totalRead == available) {
                    float progress = (totalRead * 100.0f) / available;
                    Log.d(TAG, "加载进度: " + String.format("%.1f", progress) + "%");
                }
            }
            
            if (totalRead != available) {
                throw new Exception("文件读取不完整: " + totalRead + "/" + available);
            }
            
            return modelBytes;
        }
    }
    
    /**
     * 打印模型信息
     */
    private void printModelInfo() {
        try {
            Log.d(TAG, "=== 模型信息 ===");
            Log.d(TAG, "输入节点: " + session.getInputNames().size());
            Log.d(TAG, "输出节点: " + session.getOutputNames().size());
            
            for (String name : session.getInputNames()) {
                Log.d(TAG, "输入: " + name);
            }
            
            for (String name : session.getOutputNames()) {
                Log.d(TAG, "输出: " + name);
            }
        } catch (Exception e) {
            Log.w(TAG, "获取模型信息失败", e);
        }
    }
    
    /**
     * 执行推理（简化版本）
     */
    public Bitmap removeWatermark(Bitmap originalImage, Bitmap maskImage) throws Exception {
        if (session == null || env == null) {
            throw new IllegalStateException("模型未初始化");
        }
        
        Log.d(TAG, "开始简化版推理");
        long startTime = System.currentTimeMillis();
        
        try {
            // 1. 预处理
            float[] inputData = preprocessInput(originalImage, maskImage);
            
            // 2. 创建输入张量
            long[] shape = {1, 4, originalImage.getHeight(), originalImage.getWidth()};
            OnnxTensor inputTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), shape);
            
            // 3. 执行推理
            try (OrtSession.Result result = session.run(Collections.singletonMap(INPUT_NAME, inputTensor))) {
                
                // 4. 获取输出
                Object outputObj = result.get(0).getValue();
                float[] outputData = convertOutput(outputObj);
                
                // 5. 后处理
                Bitmap resultBitmap = postprocessOutput(outputData, originalImage.getWidth(), originalImage.getHeight());
                
                long endTime = System.currentTimeMillis();
                Log.d(TAG, "简化版推理完成，耗时: " + (endTime - startTime) + "ms");
                
                return resultBitmap;
            }
            
        } catch (Exception e) {
            Log.e(TAG, "推理失败", e);
            throw e;
        }
    }
    
    /**
     * 简化的预处理
     */
    private float[] preprocessInput(Bitmap image, Bitmap mask) {
        int width = image.getWidth();
        int height = image.getHeight();
        int totalPixels = width * height;
        
        float[] inputData = new float[4 * totalPixels]; // RGBA
        
        int[] imagePixels = new int[totalPixels];
        int[] maskPixels = new int[totalPixels];
        
        image.getPixels(imagePixels, 0, width, 0, 0, width, height);
        
        // 调整mask尺寸
        Bitmap resizedMask = Bitmap.createScaledBitmap(mask, width, height, true);
        resizedMask.getPixels(maskPixels, 0, width, 0, 0, width, height);
        
        // 转换为NCHW格式
        for (int i = 0; i < totalPixels; i++) {
            int imagePixel = imagePixels[i];
            int maskPixel = maskPixels[i];
            
            // RGB通道 (归一化到[-1,1])
            inputData[i] = (Color.red(imagePixel) / 255.0f) * 2.0f - 1.0f;
            inputData[totalPixels + i] = (Color.green(imagePixel) / 255.0f) * 2.0f - 1.0f;
            inputData[2 * totalPixels + i] = (Color.blue(imagePixel) / 255.0f) * 2.0f - 1.0f;
            
            // Mask通道
            inputData[3 * totalPixels + i] = Color.red(maskPixel) / 255.0f;
        }
        
        if (resizedMask != mask) {
            resizedMask.recycle();
        }
        
        return inputData;
    }
    
    /**
     * 简化的输出转换
     */
    private float[] convertOutput(Object outputObj) {
        if (outputObj instanceof float[][][][]) {
            float[][][][] output4D = (float[][][][]) outputObj;
            int channels = output4D[0].length;
            int height = output4D[0][0].length;
            int width = output4D[0][0][0].length;
            
            float[] result = new float[channels * height * width];
            int index = 0;
            for (int c = 0; c < channels; c++) {
                for (int h = 0; h < height; h++) {
                    for (int w = 0; w < width; w++) {
                        result[index++] = output4D[0][c][h][w];
                    }
                }
            }
            return result;
        }
        throw new RuntimeException("不支持的输出格式");
    }
    
    /**
     * 简化的后处理
     */
    private Bitmap postprocessOutput(float[] outputData, int width, int height) {
        Bitmap result = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        int[] pixels = new int[width * height];
        int totalPixels = width * height;
        
        for (int i = 0; i < totalPixels; i++) {
            // 从[-1,1]转换到[0,255]
            float r = (outputData[i] + 1.0f) * 127.5f;
            float g = (outputData[totalPixels + i] + 1.0f) * 127.5f;
            float b = (outputData[2 * totalPixels + i] + 1.0f) * 127.5f;
            
            int red = Math.max(0, Math.min(255, (int) r));
            int green = Math.max(0, Math.min(255, (int) g));
            int blue = Math.max(0, Math.min(255, (int) b));
            
            pixels[i] = Color.argb(255, red, green, blue);
        }
        
        result.setPixels(pixels, 0, width, 0, 0, width, height);
        return result;
    }
    
    /**
     * 清理资源
     */
    private void cleanup() {
        try {
            if (session != null) {
                session.close();
                session = null;
            }
            if (env != null) {
                env.close();
                env = null;
            }
        } catch (Exception e) {
            Log.e(TAG, "清理资源失败", e);
        }
    }
    
    /**
     * 释放资源
     */
    public void release() {
        Log.d(TAG, "释放简化版LaMa资源");
        cleanup();
    }
    
    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return session != null && env != null;
    }
}
